# -*- coding:utf-8 -*-
import struct
import sys
import os
import re

RETURN_OK = 0
RETURN_ERR = -1
BLOCK_SIZE = 2048

HS_MAGIC = b"HSDT"
HS_VERSION = 1
HS_MATCH_KEY = r"hisi,boardid = <(.+)>"
BOARDID_CNT = 4

class StatFile:
    def __init__(self, filename, boardid, offset):
        statinfo = os.stat(filename)
        self.filename = filename
        self.boardid = boardid
        self.rsv = [0, 0, 0, 0]
        self.dtb_size = statinfo.st_size + (BLOCK_SIZE - (statinfo.st_size % BLOCK_SIZE))
        self.vrl_size = 0
        self.dtb_offset = offset
        self.vrl_offset = 0
        self.dtb_file = 0
        self.vrl_file = 0
        self.node_list = [self.filename, self.boardid, self.rsv, self.dtb_size, self.vrl_file,
                          self.dtb_offset, self.vrl_offset, self.dtb_file, self.vrl_file]

class MakeDtb:
    def __init__(self, outname, dtc_path, dtb_path):
        self.outname = outname
        self.dtb_node = []
        offset = BLOCK_SIZE
        dtb_list = self.get_dtb_files(dtb_path)
        self.dtb_cnt = len(dtb_list)
        self.ret = True
        print("Found %d DTB(s)" % (self.dtb_cnt))
        if self.dtb_cnt == 0:
            print("please check dtb path.")
            self.ret = False
            return

        # sort list
        dtb_list.sort(key = lambda x : self.calculate_boardid(x))

        # add list
        for dtb in dtb_list:
            boardid = self.match_boardid(dtc_path, dtb_path, dtb)
            if boardid != RETURN_ERR:
                self.dtb_node.append(StatFile(dtb_path + dtb, boardid, offset).node_list)
                offset += StatFile(dtb_path + dtb, boardid, offset).dtb_size
            else:
                print("boardid match failed.")
                self.ret = False
                return

        # write file
        self.write_dt()

    def calculate_boardid(self, filename):
        boardid = int(int(re.split("\.|\-", filename)[-2]) / 10)
        return boardid

    def get_dtb_files(self, dtb_path):
        dtb_list = []
        for root, dirs, files in os.walk(dtb_path):
            for file_dtb in files:
                if file_dtb.split('.')[-1] == "dtb":
                    dtb_list.append(file_dtb)
        return dtb_list

    def match_boardid(self, dtc_path, dtb_path, dtb):
        cmd = "%s/dtc -I dtb -O dts 2>&1 %s%s" % (dtc_path, dtb_path, dtb)
        with os.popen(cmd ,"r") as pf:
            line = pf.readline()
            while line:
                if re.findall(HS_MATCH_KEY, line):
                    boardid = re.findall(HS_MATCH_KEY, line)[0].split(' ')
                    return [int(i, 16) for i in boardid]
                line = pf.readline()
        return RETURN_ERR

    def write_boardid(self, cnt, fd, list):
        for i in list:
            fd.write(struct.pack("B", i))

    def write_dtb(self, fd, list):
        with open(list[0], "rb") as f_tmp:
            data = f_tmp.read()
            fd.seek(list[5])
            fd.write(data)

    def write_dt(self):
        header_len = 12 + 40 * self.dtb_cnt + 4
        filler_len = BLOCK_SIZE - (header_len % BLOCK_SIZE)

        with open(self.outname, "wb") as f:
            # write magic
            f.write(HS_MAGIC)
            f.write(struct.pack("I", HS_VERSION))
            f.write(struct.pack("I", self.dtb_cnt))

            # dt_entry (size, offset)
            for i in self.dtb_node:
                self.write_boardid(BOARDID_CNT, f, i[1]) # board id: dtb_node[i][1]
                self.write_boardid(BOARDID_CNT, f, i[2]) # reserved: dtb_node[i][2]
                f.write(struct.pack("IIIILL", i[3],      # dtb_size: dtb_node[i][3]
                                              i[4],      # vrl_size: dtb_node[i][4]
                                              i[5],      # dtb_offset: dtb_node[i][5]
                                              i[6],      # vrl_offset: dtb_node[i][6]
                                              i[7],      # dtb_offset: dtb_node[i][7]
                                              i[8]))     # vrl_offset: dtb_node[i][8]
            f.write(struct.pack("I", 0))                 # END

            # write dtb
            f.seek(header_len + filler_len)
            for i in self.dtb_node:
                self.write_dtb(f, i)

            padding = BLOCK_SIZE - (f.tell() % BLOCK_SIZE)
            if padding > 0:
                f.seek(f.tell() + padding - 1)
                f.write(struct.pack("B", 0))

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print('Usage: %s out[dt.img] dtc_path dtb_path' % (sys.argv[0]))
        sys.exit(1)

    # check para
    dtc_path = os.path.join(sys.argv[2], "dtc")
    if not os.path.exists(dtc_path) or not os.path.exists(sys.argv[3]):
        print("please check path is exist.")
        print('Usage: %s out[dt.img] dtc_path dtb_path' % (sys.argv[0]))
        sys.exit(1)

    print("Out file: %s\nDTC path: %s\nDTB path: %s" % (sys.argv[1], sys.argv[2], sys.argv[3]))
    if MakeDtb(sys.argv[1], sys.argv[2], sys.argv[3]).ret:
        sys.exit(0)
    else:
        print("mkdtb fail!")
        sys.exit(1)
